use super::*;
use riir::agent;
use riir::agent::tools::playground::{run_rustycli, VerifyResp, VerifyType};
use std::marker::PhantomData;
use tracing::info;

use serde::{Deserialize, Serialize};

pub const REFACTOR_RUST_PATH: &str = "./files/rust/refactor.md";

#[derive(Serialize, Deserialize)]
pub struct RustExpertInfo {
    pub rust_expert: String,
}

pub const RUST_EXPERT: Prompt<RustExpertInfo> = Prompt("RustExpert: [Rust Expert]", PhantomData);

#[derive(Serialize, Deserialize)]
pub struct RefactorPlanInfo {
    pub task: String,
}

pub const RUST_REFACTOR_PLAN: Prompt<RefactorPlanInfo> = Prompt(
    r#"
You are a Rust expert assistant. You can extract Rust code from the user's input, interpret it, analyze it step by step, and then refactor it.

The requirements for refactoring the code are as follows:

- Please ensure the code is fully functional and doesn't omit any details.
- The strings in the code you generate do not need extra escape characters.
- Identify syntax errors in the code and ensure that the syntax complies with stable Rust.
- Identify the syntax style in the code to meet authentic Rust usage, including semantic naming.
- Recognize the overall architecture and eliminate duplicate code.
- The refactored code should have the same functionality as the original.
- The Rust code snippet after your refactoring must not have any escape characters, please!


Please ensure that your reply strictly adheres to the following  JSON format:

```json
{{
	"plan": "<concise plan on how you will complete the task>",
	"code": "
	```rust
	<The Rust code after your refactoring>
	```
	"
}}
```
"#,
    PhantomData,
);

#[derive(Debug)]
pub struct RefactorAgent {
    attributes: AgentBase,
    bug_errors: Option<String>,
    bug_count: u8,
    retry_count: u8,
    llm: Arc<Mutex<Box<dyn LLM + 'static>>>,
}

impl RefactorAgent {
    pub fn new(llm: Arc<Mutex<Box<dyn LLM + 'static>>>) -> Self {
        let attributes: AgentBase = AgentBase {
            objective: "Refacting Rust Code".to_string(),
            position: "Refactor Agent".to_string(),
            state: AgentState::Starting,
            memory: vec![],
        };

        Self {
            attributes,
            bug_errors: None,
            bug_count: 0,
            retry_count: 0,
            llm: llm.clone(),
        }
    }

    async fn call_initial_code(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        // info!("--- Refactor Decode Code --- : {decode_code:?}");
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_REFACTOR_PLAN.0),
            format!(
                "CODE: {} ",
                factsheet.code.as_ref().unwrap_or(&"".to_string())
            ),
        );

        info!("--- Refactor Init Msg Code --- : {msg_context:?}");

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "refactor", &ai_response).await;
        factsheet.code = Some(ai_response);
    }

    async fn call_improved_code(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_REFACTOR_PLAN.0),
            format!(
                "CODE: {} ",
                factsheet.code.as_ref().unwrap_or(&"".to_string())
            ),
        );

        info!("--- Refactor improved Msg Code --- : {msg_context:?}");

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "refactor", &ai_response).await;
        factsheet.code = Some(ai_response);
    }

    async fn call_fix_code_bugs(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_FIXED_EXPERT.0),
            format!(
                "CODE: {:?} \n ERROR_BUGS: {:?} \n ",
                factsheet.code.as_ref().unwrap_or(&"".to_string()),
                self.bug_errors
            ),
        );

        info!("Refactor Agent Debug: Fix Code Msg --- {msg_context:?}");

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "refactor", &ai_response).await;
        factsheet.code = Some(ai_response);
    }
}

#[async_trait]
impl AgentFns for RefactorAgent {
    fn get_attributes_from_agent(&self) -> &AgentBase {
        &self.attributes
    }

    async fn execute(
        &mut self,
        factsheet: &mut FactSheet,
        output_path: PathBuf
    ) -> Result<(), Box<dyn std::error::Error>> {
        while self.attributes.state != AgentState::Done {
            match &self.attributes.state {
                AgentState::Starting => {
                    info!("Refactor Agent Start ---");
                    self.call_initial_code(factsheet, output_path.clone()).await;
                    self.attributes.state = AgentState::Working;
                    continue;
                }

                AgentState::Working => {
                    info!("Refactor Agent Working ---");
                    if self.bug_count == 0 {
                        self.call_improved_code(factsheet, output_path.clone()).await;
                    } else {
                        self.call_fix_code_bugs(factsheet, output_path.clone()).await;
                    }
                    self.attributes.state = AgentState::Debug;
                    continue;
                }

                AgentState::Debug => {
                    info!("Refactor Agent Debug ---");

                    // --- 验证 LLM 响应代码
                    if let Some(e) = agent::capture_rust_code(&factsheet.code.clone().unwrap()) {
                        info!("EndANswer Chat: plain {:?}", e);

                        // do something async
                        let verft = run_rustycli(&e, VerifyType::BUILD).await;
                        info!("Verify Chat: {:?}", verft);
                        match verft {
                            VerifyResp::Success { stdout, stderr } => {
                                self.bug_count = 0;
                                info!("Code build successful...");
                                self.attributes.state = AgentState::Done;
                            }
                            VerifyResp::Fail(error_str) => {
                                info!("Refactor Agent Debug: Bug Error --- {error_str}");
                                // Update error stats
                                self.bug_count += 1;
                                self.retry_count += 1;
                                self.bug_errors = Some(error_str);
                                if self.retry_count > 10 {
                                    info!(
                                        "Code build failed but retry count is greater then 10... "
                                    );
                                    self.attributes.state = AgentState::Done;
                                } else {
                                    // Pass back for rework
                                    self.attributes.state = AgentState::Working;
                                    continue;
                                }
                            }
                        }
                    } else {
                        info!("Uncapture rust code!");
                        // Pass back for rework
                        self.attributes.state = AgentState::Working;
                        continue;
                    }
                }

                _ => {}
            }
        }
        Ok(())
    }
}
